import os
import numpy as np
import cv2
import random
from PIL import Image
import torchvision.transforms as transforms


def shuffle_split_data(data,testratio=0.2):
    # random.shuffle(real)
    # random.shuffle(fake)
    # len1 = len(real)
    # len2 = len(fake)
    # tlen1 = int(len1*testratio)
    # tlen2 = int(len2*testratio)
    
    # realtest,faketest = real[:tlen1],fake[:tlen2]
    # realtrain,faketrain = real[tlen1:],fake[tlen2:]
    
    random.shuffle(data)
    len1 = len(data)
    tlen1 = int(len1*testratio)
    
    datatest= data[:tlen1]
    datatrain= data[tlen1:]

    return datatrain,datatest

def searchalldata(root):
    domains = os.listdir(root)
    domains.remove('gendata.py')
    dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            s0paths = os.listdir(os.path.join(root,domain,cat,'S0'))
            dolppaths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
            s0paths.sort()
            dolppaths.sort()
            if 'geniune' in domain:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0real.append((os.path.join(root,domain,cat,s0path),1))
                    dolpreal.append((os.path.join(root,domain,cat,dolppath),1))
            else:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0fake.append((os.path.join(root,domain,cat,s0path),0))
                    dolpfake.append((os.path.join(root,domain,cat,dolppath),0))
    return dolpreal,dolpfake,s0real,s0fake

def split_real_fake(dolpreal,dolpfake,s0real,s0fake,testratio=0.2):
    if len(dolpreal) > 0 :
        indreal = np.arange(0,len(dolpreal))
        indrealtrain,indrealtest = shuffle_split_data(indreal,testratio)
        dolpreal,s0real = np.array(dolpreal),np.array(s0real)
        
    if len(dolpfake) > 0:
        indfake = np.arange(0,len(dolpfake))
        indfaketrain,indfaketest = shuffle_split_data(indfake,testratio)
        dolpfake,s0fake = np.array(dolpfake),np.array(s0fake)
    # indrealtrain,indfaketrain,indrealtest,indfaketest = shuffle_split_data(indreal,indfake,testratio)
    
    # dolpreal,dolpfake,s0real,s0fake = np.array(dolpreal),np.array(dolpfake),np.array(s0real),np.array(s0fake)
    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    if len(dolpreal) > 0 and len(dolpfake) > 0:
        # dolptraindir,dolptestdir = np.vstack((dolpreal[indrealtrain],dolpfake[indfaketrain])),np.vstack((dolpreal[indrealtest],dolpfake[indfaketest]))
        # s0traindir,s0testdir = np.vstack((s0real[indrealtrain],s0fake[indfaketrain])),np.vstack((s0real[indrealtest],s0fake[indfaketest]))
        dolptraindir,dolptestdir = list(dolpreal[indrealtrain])+list(dolpfake[indfaketrain]),list(dolpreal[indrealtest])+list(dolpfake[indfaketest])
        s0traindir,s0testdir = list(s0real[indrealtrain])+list(s0fake[indfaketrain]),list(s0real[indrealtest])+list(s0fake[indfaketest])
    elif len(dolpreal) > 0 and len(dolpfake) == 0:
        dolptraindir,dolptestdir = list(dolpreal[indrealtrain]),list(dolpreal[indrealtest])
        if len(s0real) == len(dolpreal):
            s0traindir,s0testdir = list(s0real[indrealtrain]),list(s0real[indrealtest])
    elif len(dolpreal) == 0 and len(dolpfake) > 0:
        dolptraindir,dolptestdir = list(dolpfake[indfaketrain]),list(dolpfake[indfaketest])
        if len(s0fake) == len(dolpfake):
            s0traindir,s0testdir = list(s0fake[indfaketrain]),list(s0fake[indfaketest])
    
    # dolptraindir+=_dolptraindir
    # dolptestdir+=dolptestdir
    # s0traindir+=s0traindir
    # s0testdir+=s0testdir
        
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)

def gendatalist(root=None,testratio=0.2):    
    dolpreal,dolpfake,s0real,s0fake = searchalldata(root)
    return split_real_fake(dolpreal,dolpfake,s0real,s0fake,testratio)

def save_train_test_to_txt(traindir,testdir,txtsavepath,filenameprefix,version):
    # traindir,testdir=gendatalist(datadir)
    # with open(os.path.join(txtsavepath,filenameprefix+'all.txt'),'w') as f:
    #     all_ = traindir + testdir
    #     for line in all_:
    #         f.write(line[0]+','+str(line[1])+'\n')
            
    with open(os.path.join(txtsavepath,filenameprefix+'_'+version+'_train.txt'),'w') as f:
        for line in traindir:
            f.write(line[0]+','+str(line[1])+'\n')
    
    with open(os.path.join(txtsavepath,filenameprefix+'_'+version+'_test.txt'),'w') as f:
        for line in testdir:
            f.write(line[0]+','+str(line[1])+'\n')


def gendatalist2(root):
    # domains = ['fake']['S0']
    domains=['physical_domain']
    # domains.remove('gendata.py')
    dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            s0paths = os.listdir(os.path.join(root,domain,cat,'S0'))
            dolppaths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
            s0paths.sort()
            dolppaths.sort()
            if 'geniune' in domain:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0real.append((os.path.join(root,domain,cat,'S0',s0path),1))
                    dolpreal.append((os.path.join(root,domain,cat,'DOLP',dolppath),1))
            else:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0fake.append((os.path.join(root,domain,cat,'S0',s0path),0))
                    dolpfake.append((os.path.join(root,domain,cat,'DOLP',dolppath),0))
    return dolpreal,dolpfake,s0real,s0fake

def gendatalist3(root,testratio=0.2):
    domains =os.listdir(root)
    domains.remove('gendata.py')
    # dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    data_dict = {}
    data_dict['real']={}
    data_dict['fake']={}
    data_dict['real']['S0']={}
    data_dict['real']['DOLP']={}
    data_dict['fake']['S0']={}
    data_dict['fake']['DOLP']={}
    
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            if domain == 'geniune_domain' and cat == 'HUT':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    tmp = int(path.split('_')[0])
                    if tmp in data_dict['real']['S0'].keys():
                        data_dict['real']['S0'][tmp]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['real']['S0'][tmp]=[]
                        data_dict['real']['S0'][tmp] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                pass
            elif domain == 'physical_domain':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['S0'].keys():
                        data_dict['fake']['S0'][cat]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['fake']['S0'][cat]=[]
                        data_dict['fake']['S0'][cat] += [os.path.join(root,domain,cat,'S0',path)]
                        pass

            elif domain == 'digital_domain' and cat == 'cyclegan':
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                paths.sort()
                for path in paths:
                    tmp = int(path.split('_')[0])
                    if tmp in data_dict['real']['DOLP'].keys():
                        data_dict['real']['DOLP'][tmp]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['real']['DOLP'][tmp]=[]
                        data_dict['real']['DOLP'][tmp] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass
                    # .append((os.path.join(root,domain,cat,'DOLP',path),0))
            elif domain == 'digital_domain' and cat == 'fake_generation_v4':
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                paths.sort()
                if len(data_dict['fake']['DOLP'].keys()) == 0:
                    for tmp in data_dict['fake']['S0'].keys():
                        data_dict['fake']['DOLP'][tmp] = []
                for path in paths:
                    if 'dataset3_attack_3' in path or 'A4纸' in path:
                        data_dict['fake']['DOLP']['A4_paper'] += [os.path.join(root,domain,cat,'DOLP',path)]

                    elif 'dataset3_attack_1' in path or 'dataset3_attack_6' in path or '显示屏' in path:
                        data_dict['fake']['DOLP']['screen'] += [os.path.join(root,domain,cat,'DOLP',path)]
     
                    elif 'dataset3_attack_2' in path or '相纸' in path:
                        data_dict['fake']['DOLP']['photo'] += [os.path.join(root,domain,cat,'DOLP',path)]
                        
                    elif 'dataset3_attack_4' in path :
                        data_dict['fake']['DOLP']['silicon_head'] += [os.path.join(root,domain,cat,'DOLP',path)]
                        
                    elif 'dataset3_attack_5' in path or '面具' in path:
                        data_dict['fake']['DOLP']['silicon_mask'] += [os.path.join(root,domain,cat,'DOLP',path)]
                        # pass
                    else:
                        print(path)
                    # data_dict['real']['DOLP'].append((os.path.join(root,domain,cat,'DOLP',dolppath),1))
                
            # dolppaths.sort()
            # if 'geniune' in domain:
            #     for (s0path,dolppath) in zip(s0paths,dolppaths):
            #         s0real.append((os.path.join(root,domain,cat,'S0',s0path),1))
            #         dolpreal.append((os.path.join(root,domain,cat,'DOLP',dolppath),1))
            # else:
            #     for (s0path,dolppath) in zip(s0paths,dolppaths):
            #         s0fake.append((os.path.join(root,domain,cat,'S0',s0path),0))
            #         dolpfake.append((os.path.join(root,domain,cat,'DOLP',dolppath),0))
    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    for _type in data_dict.keys():
        if _type == 'real':
            label = 1
        else:
            label = 0
        # for modality in data_dict[_type].keys():
        
        for cat in data_dict[_type]['S0'].keys():
            
            testlen = int(len(data_dict[_type]['S0'][cat]) *testratio)
            ind = np.arange(len(data_dict[_type]['S0'][cat]))
            random.shuffle(ind)
            testinds = ind[:testlen]
            traininds = ind[testlen:]
            
            for testind in testinds:
                s0testdir.append((data_dict[_type]['S0'][cat][testind],label))
                dolptestdir.append((data_dict[_type]['DOLP'][cat][testind],label))
            
            for trainind in traininds:
                s0traindir.append((data_dict[_type]['S0'][cat][trainind],label))
                dolptraindir.append((data_dict[_type]['DOLP'][cat][trainind],label))
                            
                       
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)

def gendatalist4(root,testratio=0.2):
    domains =os.listdir(root)
    domains.remove('gendata.py')
    # dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    data_dict = {}
    data_dict['real']={}
    data_dict['fake']={}
    data_dict['real']['S0']={}
    data_dict['real']['DOLP']={}
    data_dict['fake']['S0']={}
    data_dict['fake']['DOLP']={}
    
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            if domain == 'geniune_domain' and cat == 'HUT':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                        
                    if tmp in data_dict['real']['S0'].keys():
                        data_dict['real']['S0'][tmp]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['real']['S0'][tmp]=[]
                        data_dict['real']['S0'][tmp] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                    if tmp in data_dict['real']['DOLP'].keys():
                        data_dict['real']['DOLP'][tmp]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['real']['DOLP'][tmp]=[]
                        data_dict['real']['DOLP'][tmp] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass 

            elif domain == 'physical_domain':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['S0'].keys():
                        data_dict['fake']['S0'][cat]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['fake']['S0'][cat]=[]
                        data_dict['fake']['S0'][cat] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['DOLP'].keys():
                        data_dict['fake']['DOLP'][cat]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['fake']['DOLP'][cat]=[]
                        data_dict['fake']['DOLP'][cat] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass

    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    for _type in data_dict.keys():
        if _type == 'real':
            label = 1
        else:
            label = 0
        # for modality in data_dict[_type].keys():
        
        for cat in data_dict[_type]['S0'].keys():
            
            testlen = int(len(data_dict[_type]['S0'][cat]) *testratio)
            ind = np.arange(len(data_dict[_type]['S0'][cat]))
            random.shuffle(ind)
            testinds = ind[:testlen]
            traininds = ind[testlen:]
            
            for testind in testinds:
                s0testdir.append((data_dict[_type]['S0'][cat][testind],label))
                dolptestdir.append((data_dict[_type]['DOLP'][cat][testind],label))
            
            for trainind in traininds:
                s0traindir.append((data_dict[_type]['S0'][cat][trainind],label))
                dolptraindir.append((data_dict[_type]['DOLP'][cat][trainind],label))
                            
                       
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)

def gendatalist5(root,testratio=0.2):
    domains =os.listdir(root)
    domains.remove('gendata.py')
    # dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    data_dict = {}
    data_dict['real']={}
    data_dict['fake']={}
    data_dict['real']['S0']={}
    data_dict['real']['DOLP']={}
    data_dict['fake']['S0']={}
    data_dict['fake']['DOLP']={}
    
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            if domain == 'geniune_domain' :# and cat == 'HUT':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                        
                    if tmp in data_dict['real']['S0'].keys():
                        data_dict['real']['S0'][tmp]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['real']['S0'][tmp]=[]
                        data_dict['real']['S0'][tmp] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                        
                    if tmp in data_dict['real']['DOLP'].keys():
                        data_dict['real']['DOLP'][tmp]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['real']['DOLP'][tmp]=[]
                        data_dict['real']['DOLP'][tmp] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass 
                    
                pass
            elif domain == 'physical_domain':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['S0'].keys():
                        data_dict['fake']['S0'][cat]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['fake']['S0'][cat]=[]
                        data_dict['fake']['S0'][cat] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['DOLP'].keys():
                        data_dict['fake']['DOLP'][cat]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['fake']['DOLP'][cat]=[]
                        data_dict['fake']['DOLP'][cat] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass
                    
            elif domain == 'digital_domain' and cat == 'cyclegan':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['S0'].keys():
                        data_dict['fake']['S0'][cat]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['fake']['S0'][cat]=[]
                        data_dict['fake']['S0'][cat] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['DOLP'].keys():
                        data_dict['fake']['DOLP'][cat]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['fake']['DOLP'][cat]=[]
                        data_dict['fake']['DOLP'][cat] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass

    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    for _type in data_dict.keys():
        if _type == 'real':
            label = 1
        else:
            label = 0
        # for modality in data_dict[_type].keys():
        
        for cat in data_dict[_type]['S0'].keys():
            
            testlen = int(len(data_dict[_type]['S0'][cat]) *testratio)
            ind = np.arange(len(data_dict[_type]['S0'][cat]))
            random.shuffle(ind)
            testinds = ind[:testlen]
            traininds = ind[testlen:]
            
            for testind in testinds:
                s0testdir.append((data_dict[_type]['S0'][cat][testind],label))
                dolptestdir.append((data_dict[_type]['DOLP'][cat][testind],label))
            
            for trainind in traininds:
                s0traindir.append((data_dict[_type]['S0'][cat][trainind],label))
                dolptraindir.append((data_dict[_type]['DOLP'][cat][trainind],label))
                            
                       
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)

def gendatalist6(root,testratio=0.2):
    # domains =os.listdir(root)
    # domains.remove('gendata.py')
    domains = ['digital_domain']
    cats = ['cyclegan_fft','fake_generation_v4']
    # dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    data_dict = {}
    data_dict['real']={}
    data_dict['fake']={}
    data_dict['real']['S0']={}
    data_dict['real']['DOLP']={}
    data_dict['fake']['S0']={}
    data_dict['fake']['DOLP']={}
    
    for domain in domains:
        # cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            if cat == 'cyclegan_fft':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                        
                    if tmp in data_dict['real']['S0'].keys():
                        data_dict['real']['S0'][tmp]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['real']['S0'][tmp]=[]
                        data_dict['real']['S0'][tmp] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    # data_dict['real']['S0'].append((os.path.join(root,domain,cat,'S0',path),0))
                    if '真人' in path:
                        tmp = path.split('_')[0]
                    else:
                        tmp = int(path.split('_')[0])
                    if tmp in data_dict['real']['DOLP'].keys():
                        data_dict['real']['DOLP'][tmp]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['real']['DOLP'][tmp]=[]
                        data_dict['real']['DOLP'][tmp] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass 

            elif cat == 'fake_generation_v4':
                paths = os.listdir(os.path.join(root,domain,cat,'S0'))
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['S0'].keys():
                        data_dict['fake']['S0'][cat]+= [os.path.join(root,domain,cat,'S0',path)]
                    else:
                        data_dict['fake']['S0'][cat]=[]
                        data_dict['fake']['S0'][cat] += [os.path.join(root,domain,cat,'S0',path)]
                        pass
                x1 = len(paths)  
                paths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
                x2 = len(paths)
                if x1 != x2:
                    print('ERROR Mismatch')
                paths.sort()
                for path in paths:
                    if cat in data_dict['fake']['DOLP'].keys():
                        data_dict['fake']['DOLP'][cat]+= [os.path.join(root,domain,cat,'DOLP',path)]
                    else:
                        data_dict['fake']['DOLP'][cat]=[]
                        data_dict['fake']['DOLP'][cat] += [os.path.join(root,domain,cat,'DOLP',path)]
                        pass

    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    for _type in data_dict.keys():
        if _type == 'real':
            label = 1
        else:
            label = 0
        # for modality in data_dict[_type].keys():
        
        for cat in data_dict[_type]['S0'].keys():
            
            testlen = int(len(data_dict[_type]['S0'][cat]) *testratio)
            ind = np.arange(len(data_dict[_type]['S0'][cat]))
            random.shuffle(ind)
            testinds = ind[:testlen]
            traininds = ind[testlen:]
            
            for testind in testinds:
                s0testdir.append((data_dict[_type]['S0'][cat][testind],label))
                dolptestdir.append((data_dict[_type]['DOLP'][cat][testind],label))
            
            for trainind in traininds:
                s0traindir.append((data_dict[_type]['S0'][cat][trainind],label))
                dolptraindir.append((data_dict[_type]['DOLP'][cat][trainind],label))
                            
                       
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)


def initdirs(dirs):
    if os.path.exists(dirs[0]) is False:
        for _dir in dirs:
            os.makedirs(_dir)
    else:
        def del_dirsfile(dirs):
            for _dir in dirs:
                files = os.listdir(_dir)
                for file in files:
                    os.remove(os.path.join(_dir,file))
        del_dirsfile(dirs)

if __name__ == '__main__':
    datadir = os.path.join(os.path.abspath('.'),'datasets','CDFA')
    
    mode = 'gen6'
    
    if mode == 'gen1':
        ratio = 0.99
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','DataRatio_'+version)
        initdirs([txtdir])
            
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist(datadir,ratio)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0testdir,txtdir,'S0_',version)     
    # elif mode == 'test':
    #     m = MultiPolarFaceDataset(datadir,'train')
    #     m.out='multi'
    #     m.__getitem__(150)
    elif mode == 'gen2':
        ratio = 0
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','fake_generation')
        initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist2(datadir)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptestdir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0testdir,s0testdir,txtdir,'S0_',version) 
    elif mode == 'gen2.5':
        ratio = 0
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','real_generation')
        initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist2(datadir)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptraindir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0traindir,txtdir,'S0_',version) 
    elif mode == 'gen3':
        ratio = 0.90
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','CDCN_2modality_v4')
        # initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist3(datadir,ratio)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0testdir,txtdir,'S0_',version)
    
    elif mode == 'gen4':
        ratio = 0.99
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','CDFA_LD_No_digital')
        # initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist4(datadir,ratio)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0testdir,txtdir,'S0_',version) 
    
    elif mode == 'gen5':
        ratio = 0.99
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','CDFA')
        # initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist5(datadir,ratio)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0testdir,txtdir,'S0_',version)
    
    elif mode == 'gen6':
        ratio = 0.99
        version = 'v'+str(int(100*ratio))
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207071430','CDFA_Only_Digital')
        # initdirs([txtdir])
        
        dolptraindir,dolptestdir,s0traindir,s0testdir = gendatalist6(datadir,ratio)
        print('DOLP:  trainLen:'+str(len(dolptraindir))+'  testLen:'+str(len(dolptestdir)))
        print('S0:  trainLen:'+str(len(s0traindir))+'  testLen:'+str(len(s0testdir)))
        save_train_test_to_txt(dolptraindir,dolptestdir,txtdir,'DOLP_',version)
        save_train_test_to_txt(s0traindir,s0testdir,txtdir,'S0_',version) 
    pass